import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_sparse import SparseTensor


def xavier_uniform_(tensor, gain=1.):
    # tensor.size()返回tensor的尺寸（形状）。[-2:]表示取最后两个维度的大小，即fan_in（输入神经元数）和fan_out（输出神经元数）。
    # 这是在全连接层（线性层）中使用的，通常代表输入和输出的特征维度。
    fan_in, fan_out = tensor.size()[-2:]
    # std是标准差，基于输入和输出神经元数量计算。gain是一个可调参数，用于放大或缩小标准差。公式math.sqrt(2.0 / float(fan_in + fan_out))是Xavier初始化中的标准差计算公式，用于保持权重的方差在传播过程中稳定，避免梯度消失或爆炸。
    std = gain * math.sqrt(2.0 / float(fan_in + fan_out))
    # Xavier初始化使用均匀分布来初始化权重，范围是[-a, a]。这里a是根据标准差std计算出来的，用公式a = sqrt(3) * std。这个公式确保均匀分布的方差等于std的方差。
    a = math.sqrt(3.0) * std  # Calculate uniform bounds from standard deviation
    # torch.nn.init._no_grad_uniform_是PyTorch内部函数，用于在不计算梯度的情况下，将tensor中的元素初始化为均匀分布，范围是[-a, a]。这个初始化方法有助于保持网络训练时的梯度稳定。
    return torch.nn.init._no_grad_uniform_(tensor, -a, a)


class Transformer(nn.Module):
    '''
        The transformer-based semantic fusion in SeHGNN.
    '''
    def __init__(self, n_channels, num_heads=1, att_drop=0., act='none'):
        super(Transformer, self).__init__()
        # 存储输入数据的通道数（特征维度）。
        self.n_channels = n_channels
        # 存储多头注意力机制的头数。
        self.num_heads = num_heads
        # 确保通道数能够均匀地分配到多头注意力机制中，并且每个头有合适的维度进行计算（注意力头的维度是通道数的1/4）。
        assert self.n_channels % (self.num_heads * 4) == 0

        # 定义一个线性层，将输入数据从n_channels维度映射到n_channels//4维度，用于计算查询矩阵。
        self.query = nn.Linear(self.n_channels, self.n_channels//4)
        # 定义一个线性层，将输入数据映射到n_channels//4维度，用于计算键矩阵。
        self.key   = nn.Linear(self.n_channels, self.n_channels//4)
        # 定义一个线性层，将输入数据映射回n_channels维度，用于计算值矩阵。
        self.value = nn.Linear(self.n_channels, self.n_channels)

        # 定义一个可学习的标量参数gamma，用于在最终的输出中调整注意力的影响。
        self.gamma = nn.Parameter(torch.tensor([0.]))
        # 定义一个Dropout层，用于在计算注意力矩阵时进行随机失活，防止过拟合。
        self.att_drop = nn.Dropout(att_drop)
        # 根据act参数选择不同的激活函数，并将其保存到self.act中。如果指定的激活函数未被识别，会触发assert语句抛出错误。
        if act == 'sigmoid':
            self.act = torch.nn.Sigmoid()
        elif act == 'relu':
            self.act = torch.nn.ReLU()
        elif act == 'leaky_relu':
            self.act = torch.nn.LeakyReLU(0.2)
        elif act == 'none':
            self.act = lambda x: x
        else:
            assert 0, f'Unrecognized activation function {act} for class Transformer'

        # 初始化模型中的可学习参数。
        self.reset_parameters()

    def reset_parameters(self):
        # 遍历模型中的所有子模块，并调用它们的reset_parameters方法（如果存在），以重新初始化参数。
        for k, v in self._modules.items():
            if hasattr(v, 'reset_parameters'):
                v.reset_parameters()
        # 将gamma参数初始化为零。
        nn.init.zeros_(self.gamma)

    def forward(self, x, mask=None):
        # 获取输入张量x的尺寸，分别表示批次大小B、元路径的数量M和通道数C。
        B, M, C = x.size() # batchsize, num_metapaths, channels
        # 将多头注意力机制的头数存储在H中。
        H = self.num_heads
        # 如果提供了mask（用于遮蔽某些位置的张量），确保它的尺寸与x的批次大小和元路径数量匹配。
        if mask is not None:
            assert mask.size() == torch.Size((B, M))

        # 将输入x通过query线性层进行投影，得到查询矩阵f，然后将其形状调整为[B, H, M, -1]，表示多头注意力的每个头的查询矩阵。
        f = self.query(x).view(B, M, H, -1).permute(0,2,1,3) # [B, H, M, -1]
        # 类似于f的计算，将输入x通过key线性层进行投影，得到键矩阵g，并调整形状为[B, H, -1, M]。
        g = self.key(x).view(B, M, H, -1).permute(0,2,3,1)   # [B, H, -1, M]
        # 将输入x通过value线性层进行投影，得到值矩阵h，并调整形状为[B, H, M, -1]。
        h = self.value(x).view(B, M, H, -1).permute(0,2,1,3) # [B, H, M, -1]

        # 计算注意力矩阵beta。首先，将查询矩阵f与键矩阵g相乘，并除以f的最后一维的平方根进行缩放（避免梯度消失或爆炸）。然后，应用指定的激活函数self.act，最后通过softmax函数将其归一化。
        beta = F.softmax(self.act(f @ g / math.sqrt(f.size(-1))), dim=-1) # [B, H, M, M(normalized)]
        # 对注意力矩阵beta应用Dropout，以随机丢弃一些注意力权重，防止过拟合。
        beta = self.att_drop(beta)
        # 如果提供了mask，则将其应用到注意力矩阵beta上，使得某些位置的注意力权重为零，并重新归一化。
        if mask is not None:
            beta = beta * mask.view(B, 1, 1, M)
            beta = beta / (beta.sum(-1, keepdim=True) + 1e-12)

        # 计算输出o。将注意力矩阵beta与值矩阵h相乘，得到加权后的值矩阵。然后将结果乘以标量参数gamma，用于控制注意力的强度。
        o = self.gamma * (beta @ h) # [B, H, M, -1]
        # 调整输出o的形状，使其与输入x的形状相同，并与输入x相加（即残差连接），最后返回结果。这种残差连接有助于缓解梯度消失问题。
        return o.permute(0,2,1,3).reshape((B, M, C)) + x


class LinearPerMetapath(nn.Module):
    '''
        Linear projection per metapath for feature projection in SeHGNN.
    '''
    def __init__(self, cin, cout, num_metapaths):
        super(LinearPerMetapath, self).__init__()
        # 将输入的特征维度（cin）存储到实例变量self.cin中。
        self.cin = cin
        # 将输出的特征维度（cout）存储到实例变量self.cout中。
        self.cout = cout
        # 将元路径（metapath）的数量存储到实例变量self.num_metapaths中。
        self.num_metapaths = num_metapaths

        # 定义权重矩阵W，形状为(num_metapaths, cin, cout)，表示为每个元路径都有一个独立的线性变换矩阵。
        # torch.randn用于初始化W，其中每个元素的值从标准正态分布中采样。nn.Parameter将其声明为一个可训练的参数。
        self.W = nn.Parameter(torch.randn(self.num_metapaths, self.cin, self.cout))
        # 定义偏置矩阵bias，形状为(num_metapaths, cout)，表示每个元路径都有独立的偏置。偏置矩阵的初始值为0。nn.Parameter将其声明为一个可训练的参数。
        self.bias = nn.Parameter(torch.zeros(self.num_metapaths, self.cout))

        # 调用reset_parameters方法来初始化W和bias参数，确保它们以合理的值开始训练。
        self.reset_parameters()

    def reset_parameters(self):
        # 计算适用于ReLU激活函数的增益（gain）值。增益值用于调整权重初始化的范围，确保训练的稳定性和收敛性。
        gain = nn.init.calculate_gain("relu")
        # 使用Xavier均匀分布来初始化权重矩阵W，以gain值为标准差调整初始化范围。这种初始化方式旨在保持前向和反向传播时梯度的稳定。
        xavier_uniform_(self.W, gain=gain)
        # 将偏置矩阵bias的所有元素初始化为0。
        nn.init.zeros_(self.bias)

    def forward(self, x):
        # torch.einsum('bcm,cmn->bcn', x, self.W)：使用Einstein求和约定（einsum）来实现矩阵乘法，其中x的形状为(B, C, M)，W的形状为(C, M, N)。
        # 这里bcm、cmn和bcn分别对应批量大小、输入特征维度、输出特征维度和元路径数量。这一步实现了对输入特征的线性变换。
        # self.bias.unsqueeze(0)：将bias的第一个维度扩展以匹配批量大小，然后将其添加到矩阵乘法的结果中，应用偏置。
        return torch.einsum('bcm,cmn->bcn', x, self.W) + self.bias.unsqueeze(0)


# 这行代码定义了一个匿名函数（lambda 函数），用来将嵌套列表展平（也就是将多个子列表合并成一个单一的列表）。
unfold_nested_list = lambda x: sum(x, [])

class SeHGNN(nn.Module):
    '''
        The SeHGNN model.
    '''
    def __init__(self, dataset, nfeat, hidden, nclass, feat_keys, label_feat_keys, tgt_type,
                 dropout, input_drop, att_drop, n_fp_layers, n_task_layers, act,
                 residual=False, data_size=None, drop_metapath=0., num_heads=1):
        super(SeHGNN, self).__init__()
        # 将数据集名称存储在实例变量self.dataset中，用于后续的条件处理。
        self.dataset = dataset
        # 将特征键排序，并存储在self.feat_keys中。特征键用于在特征字典中访问不同元路径的特征。
        self.feat_keys = sorted(feat_keys)
        # 将标签特征键排序，并存储在self.label_feat_keys中。这些键用于访问标签相关的特征。
        self.label_feat_keys = sorted(label_feat_keys)
        # 计算并存储模型中的通道数量，即特征键和标签特征键的总数。
        self.num_channels = num_channels = len(self.feat_keys) + len(self.label_feat_keys)
        # 存储目标节点的类型，这决定了模型要对哪个节点类型进行预测。
        self.tgt_type = tgt_type
        # 存储是否使用残差连接的布尔值。
        self.residual = residual

        # 初始化输入特征的dropout层，防止过拟合。
        self.input_drop = nn.Dropout(input_drop)

        # 存储数据尺寸信息，用于初始化嵌入参数。
        self.data_size = data_size
        # 初始化一个参数字典self.embeding，用于存储每个元路径类型的嵌入矩阵。
        self.embeding = nn.ParameterDict({})
        # 为每个元路径类型创建一个嵌入矩阵，并将其添加到self.embeding字典中。矩阵的形状为(v, nfeat)，其中v是该元路径类型的节点数量，nfeat是输入特征维度。
        for k, v in data_size.items():
            self.embeding[str(k)] = nn.Parameter(torch.Tensor(v, nfeat))

        # 检查是否存在标签特征键，如果存在，执行以下步骤。
        if len(self.label_feat_keys):
            # 初始化另一个参数字典self.labels_embeding，用于存储标签特征的嵌入矩阵。
            self.labels_embeding = nn.ParameterDict({})
            # 为每个标签特征键创建一个嵌入矩阵，并将其添加到self.labels_embeding字典中。矩阵的形状为(nclass, nfeat)，其中nclass是类别数量。
            for k in self.label_feat_keys:
                self.labels_embeding[k] = nn.Parameter(torch.Tensor(nclass, nfeat))
        else:
            self.labels_embeding = {}

        # 定义特征投影模块，包括线性层、归一化层、激活函数和dropout层，处理输入特征。
        self.feature_projection = nn.Sequential(
            *([LinearPerMetapath(nfeat, hidden, num_channels),
               nn.LayerNorm([num_channels, hidden]),
               nn.PReLU(),
               nn.Dropout(dropout),]
            + unfold_nested_list([[
               LinearPerMetapath(hidden, hidden, num_channels),
               nn.LayerNorm([num_channels, hidden]),
               nn.PReLU(),
               nn.Dropout(dropout),] for _ in range(n_fp_layers - 1)])
            )
        )

        # 初始化Transformer模块，用于进行语义融合。它接收特征投影的输出并融合不同元路径的特征。
        self.semantic_fusion = Transformer(hidden, num_heads=num_heads, att_drop=att_drop, act=act)
        # 定义一个线性层，用于将Transformer输出的特征维度从num_channels * hidden降至hidden。
        self.fc_after_concat = nn.Linear(num_channels * hidden, hidden)

        # 检查是否使用残差连接，如果是，则执行以下步骤。
        if self.residual:
            # 初始化一个线性层self.res_fc，用于将输入特征通过残差连接到最终输出。
            self.res_fc = nn.Linear(nfeat, hidden)

        # 检查当前数据集是否是IMDB或Freebase，如果不是，则初始化任务相关的MLP模块。
        if self.dataset not in ['IMDB', 'Freebase']:
            # 定义一个顺序模型self.task_mlp，包括多个线性层、批归一化层、激活函数和dropout层，用于分类任务。
            self.task_mlp = nn.Sequential(
                *([nn.PReLU(),
                   nn.Dropout(dropout),]
                + unfold_nested_list([[
                   nn.Linear(hidden, hidden),
                   nn.BatchNorm1d(hidden, affine=False),
                   nn.PReLU(),
                   nn.Dropout(dropout),] for _ in range(n_task_layers - 1)])
                + [nn.Linear(hidden, nclass),
                   nn.BatchNorm1d(nclass, affine=False, track_running_stats=False)]
                )
            )
        else:
            # 如果数据集是IMDB或Freebase，则初始化任务相关的MLP模块，但使用nn.ModuleList和不同的归一化层。
            self.task_mlp = nn.ModuleList(
                [nn.Sequential(
                    nn.PReLU(),
                    nn.Dropout(dropout))]
                + [nn.Sequential(
                    nn.Linear(hidden, hidden),
                    nn.BatchNorm1d(hidden, affine=False),
                    nn.PReLU(),
                    nn.Dropout(dropout)) for _ in range(n_task_layers - 1)]
                + [nn.Sequential(
                    nn.Linear(hidden, nclass),
                    nn.LayerNorm(nclass, elementwise_affine=False),
                    )]
            )

        # 调用reset_parameters方法初始化所有参数。
        self.reset_parameters()

    def reset_parameters(self):
        # 遍历模型中的所有模块。
        for k, v in self._modules.items():
            # 检查模块是否是参数字典类型。
            if isinstance(v, nn.ParameterDict):
                # 对参数字典中的每个参数，初始化为均匀分布在[-0.5, 0.5]之间的值。
                for _k, _v in v.items():
                    _v.data.uniform_(-0.5, 0.5)
            # 检查模块是否是模块列表类型。
            elif isinstance(v, nn.ModuleList):
                for block in v:
                    # 如果模块是顺序模型，遍历其中的每一层，并调用reset_parameters方法。
                    if isinstance(block, nn.Sequential):
                        for layer in block:
                            if hasattr(layer, 'reset_parameters'):
                                layer.reset_parameters()
                    # 如果模块具有reset_parameters方法，直接调用它。
                    elif hasattr(block, 'reset_parameters'):
                        block.reset_parameters()
            # 如果模块是顺序模型，遍历其中的每一层，并调用reset_parameters方法。
            elif isinstance(v, nn.Sequential):
                for layer in v:
                    if hasattr(layer, 'reset_parameters'):
                        layer.reset_parameters()
            # 如果模块具有reset_parameters方法，直接调用它。
            elif hasattr(v, 'reset_parameters'):
                v.reset_parameters()

    # 定义forward方法，接收一个批次的数据batch，特征字典feature_dict，标签字典label_dict，和一个可选的mask参数。
    def forward(self, batch, feature_dict, label_dict={}, mask=None):
        # 检查目标类型的特征是否是普通张量。
        if isinstance(feature_dict[self.tgt_type], torch.Tensor):
            # 对于每个元路径特征，将特征与嵌入矩阵相乘，然后应用dropout。
            features = {k: self.input_drop(x @ self.embeding[k]) for k, x in feature_dict.items()}
        # 检查目标类型的特征是否是稀疏张量。
        elif isinstance(feature_dict[self.tgt_type], SparseTensor):
            # Freebase has so many metapaths that we use feature projection per target node type instead of per metapath
            # 对稀疏张量特征执行类似的操作，但使用不同的嵌入矩阵。
            features = {k: self.input_drop(x @ self.embeding[k[-1]]) for k, x in feature_dict.items()}
        else:
            assert 0

        # 获取目标节点类型的节点数量。
        B = num_node = features[self.tgt_type].shape[0]
        # 获取通道数。
        C = self.num_channels
        # 获取特征维度。
        D = features[self.tgt_type].shape[1]

        # 对于标签特征，执行类似的dropout和嵌入矩阵相乘操作。
        labels = {k: self.input_drop(x @ self.labels_embeding[k]) for k, x in label_dict.items()}

        # 将特征和标签特征列表合并。
        x = [features[k] for k in self.feat_keys] + [labels[k] for k in self.label_feat_keys]
        # 将特征和标签特征堆叠成一个张量，形状为[B, C, D]。
        x = torch.stack(x, dim=1) # [B, C, D]
        # 通过特征投影层处理特征张量。
        x = self.feature_projection(x)

        # 通过Transformer模块进行语义融合，并转置张量以适应下一层。
        x = self.semantic_fusion(x, mask=None).transpose(1,2)

        # 通过线性层将多通道特征降维。
        x = self.fc_after_concat(x.reshape(B, -1))
        # 如果使用残差连接，则将降维后的特征与残差连接后的特征相加。
        if self.residual:
            x = x + self.res_fc(features[self.tgt_type])

        # 如果数据集不是IMDB或Freebase，直接通过MLP模块进行分类。
        if self.dataset not in ['IMDB', 'Freebase']:
            return self.task_mlp(x)
        else:
            # 如果数据集是IMDB或Freebase，使用不同的MLP结构。
            # 通过MLP的第一层。
            x = self.task_mlp[0](x)
            # 通过MLP的中间层，并添加跳跃连接。
            for i in range(1, len(self.task_mlp)-1):
                x = self.task_mlp[i](x) + x
            # 通过MLP的最后一层。
            x = self.task_mlp[-1](x)
            # 返回最终的模型输出。
            return x
